#include "net/tcp.h"
#include "net/ip.h"
#include "drivers/pcap_adapter.h"
#include "net/arp.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

#ifdef _WIN32
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#include <winsock2.h>
#include <Ws2tcpip.h>
#include <windows.h>
#else
#include <arpa/inet.h>
#include <sys/time.h>
#endif

#include <time.h>

#define MAX_CONN 32
#define TCP_RTO_MS 500
#define TCP_MAX_RETX 5   // reduced from 5 to 3 to limit SYN retransmits during tests

#pragma pack(push,1)
typedef struct {
	uint16_t src_port;
	uint16_t dst_port;
	uint32_t seq;
	uint32_t ack;
	uint8_t data_offset_reserved;
	uint8_t flags;
	uint16_t window;
	uint16_t checksum;
	uint16_t urgptr;
} tcp_hdr_t;
#pragma pack(pop)

typedef enum {CLOSED, SYN_SENT, SYN_RECEIVED, ESTABLISHED, FIN_WAIT1, FIN_WAIT2, CLOSE_WAIT, LAST_ACK, TIME_WAIT} tcp_state_t;

typedef struct {
	int used;
	uint32_t local_ip;
	uint32_t remote_ip;
	uint16_t local_port;
	uint16_t remote_port;
	uint32_t snd_nxt;
	uint32_t rcv_nxt;
	tcp_state_t state;
	uint32_t last_tx_time_ms;
	int retransmit_count;
	uint8_t pending_seg[1500];
	size_t pending_len;
} tcp_conn_t;

static tcp_conn_t g_conns[MAX_CONN];

static uint32_t now_ms() {
#ifdef _WIN32
	return (uint32_t)GetTickCount();
#else
	struct timeval tv; gettimeofday(&tv, NULL); return (uint32_t)(tv.tv_sec*1000 + tv.tv_usec/1000);
#endif
}

int tcp_init(void) {
	memset(g_conns,0,sizeof(g_conns));
	return 0;
}

static tcp_conn_t* find_conn(uint32_t r_ip, uint16_t r_port, uint16_t l_port) {
	for (int i=0;i<MAX_CONN;i++) {
		if (g_conns[i].used && g_conns[i].remote_ip==r_ip && g_conns[i].remote_port==r_port && g_conns[i].local_port==l_port) return &g_conns[i];
	}
	return NULL;
}

static tcp_conn_t* alloc_conn(void) {
	for (int i=0;i<MAX_CONN;i++) if (!g_conns[i].used) { g_conns[i].used=1; g_conns[i].state=CLOSED; return &g_conns[i]; }
	return NULL;
}

static int send_tcp_segment_via_ip(uint32_t dst_ip, uint16_t dst_port, uint16_t src_port, const uint8_t* seg, size_t seg_len) {
	return ip_send_tcp(dst_ip, dst_port, src_port, seg, seg_len);
}

static void build_and_send_tcp(tcp_conn_t* c, uint8_t flags, const uint8_t* payload, size_t payload_len) {
	tcp_hdr_t th;
	memset(&th,0,sizeof(th));
	th.src_port = htons(c->local_port);
	th.dst_port = htons(c->remote_port);
	th.seq = htonl(c->snd_nxt);
	th.ack = htonl(c->rcv_nxt);
	th.data_offset_reserved = (sizeof(tcp_hdr_t)/4)<<4;
	th.flags = flags;
	th.window = htons(1024);
	// build segment
	size_t seg_len = sizeof(tcp_hdr_t) + payload_len;
	uint8_t* seg = (uint8_t*)malloc(seg_len);
	memcpy(seg, &th, sizeof(tcp_hdr_t));
	if (payload_len) memcpy(seg + sizeof(tcp_hdr_t), payload, payload_len);
	// record for retransmission
	if (seg_len <= sizeof(c->pending_seg)) {
		memcpy(c->pending_seg, seg, seg_len);
		c->pending_len = seg_len;
		c->last_tx_time_ms = now_ms();
		c->retransmit_count = 0;
	}
	send_tcp_segment_via_ip(c->remote_ip, c->remote_port, c->local_port, seg, seg_len);
	free(seg);
}

int tcp_connect(uint32_t dst_ip, uint16_t dst_port, uint16_t src_port) {
	tcp_conn_t* c = alloc_conn();
	if (!c) return -1;
	c->remote_ip = dst_ip; c->remote_port = dst_port; c->local_port = src_port;
	c->snd_nxt = rand() & 0xFFFFFF;
	c->rcv_nxt = 0;
	c->state = SYN_SENT;
	build_and_send_tcp(c, 0x02, NULL, 0); // SYN
	return 0;
}

int tcp_send(uint32_t dst_ip, uint16_t dst_port, uint16_t src_port, const uint8_t* data, size_t len) {
	tcp_conn_t* c = find_conn(dst_ip, dst_port, src_port);
	if (!c || c->state != ESTABLISHED) return -1;
	build_and_send_tcp(c, 0x18, data, len); // PSH+ACK
	c->snd_nxt += (uint32_t)len;
	return 0;
}

int tcp_close(uint32_t dst_ip, uint16_t dst_port, uint16_t src_port) {
	tcp_conn_t* c = find_conn(dst_ip, dst_port, src_port);
	if (!c) return -1;
	if (c->state == ESTABLISHED) {
		c->state = FIN_WAIT1;
		build_and_send_tcp(c, 0x01, NULL, 0); // FIN
	}
	return 0;
}

int tcp_input(const uint8_t* tcp_pkt, size_t len, uint32_t src_ip) {
	if (len < sizeof(tcp_hdr_t)) return -1;
	const tcp_hdr_t* th = (const tcp_hdr_t*)tcp_pkt;
	uint16_t src_port = ntohs(th->src_port);
	uint16_t dst_port = ntohs(th->dst_port);
	uint32_t seq = ntohl(th->seq);
	uint32_t ack = ntohl(th->ack);
	uint8_t flags = th->flags;
	tcp_conn_t* c = find_conn(src_ip, src_port, dst_port);
	if (!c) {
		// Passive open? for PoC, respond SYN-ACK if received SYN
		if (flags & 0x02) {
			// allocate new conn
			tcp_conn_t* nc = alloc_conn();
			if (!nc) return -1;
			nc->remote_ip = src_ip; nc->remote_port = src_port; nc->local_port = dst_port;
			nc->rcv_nxt = seq + 1;
			nc->snd_nxt = rand() & 0xFFFFFF;
			nc->state = SYN_RECEIVED;
			// send SYN+ACK (SYN+ACK bits 0x12)
			build_and_send_tcp(nc, 0x12, NULL, 0);
			return 0;
		}
		return 0;
	}
	// state machine (simplified)
	if (c->state == SYN_SENT) {
		if ((flags & 0x12) == 0x12) { // SYN+ACK
			c->rcv_nxt = seq + 1;
			// send ACK
			build_and_send_tcp(c, 0x10, NULL, 0);
			c->state = ESTABLISHED;
		}
	} else if (c->state == SYN_RECEIVED) {
		if (flags & 0x10) { // ACK
			c->state = ESTABLISHED;
		}
	} else if (c->state == ESTABLISHED) {
		if (flags & 0x01) { // FIN
			c->rcv_nxt = seq + 1;
			// send ACK
			build_and_send_tcp(c, 0x10, NULL, 0);
			c->state = CLOSE_WAIT;
		}
		// If PSH or data present: for PoC we ignore payload but advance rcv_nxt and send ACK
		size_t tcp_hdr_len = (th->data_offset_reserved >> 4) * 4;
		size_t payload_len = len - tcp_hdr_len;
		if (payload_len > 0) {
			c->rcv_nxt = seq + (uint32_t)payload_len;
			build_and_send_tcp(c, 0x10, NULL, 0); // ACK
		}
	} else if (c->state == FIN_WAIT1) {
		if (flags & 0x10) { c->state = FIN_WAIT2; }
		if (flags & 0x01) {
			// remote FIN
			build_and_send_tcp(c, 0x10, NULL, 0);
			c->state = TIME_WAIT;
		}
	} else if (c->state == CLOSE_WAIT) {
		// app should call tcp_close to send FIN -> LAST_ACK
	}
	return 0;
}

void tcp_tick(uint32_t ms) {
	uint32_t now = now_ms();
	for (int i=0;i<MAX_CONN;i++) {
		if (!g_conns[i].used) continue;
		tcp_conn_t* c = &g_conns[i];
		if (c->pending_len > 0) {
			if ((now - c->last_tx_time_ms) > TCP_RTO_MS) {
				if (c->retransmit_count >= TCP_MAX_RETX) {
					// give up: cleanup
					c->used = 0;
					continue;
				}
				// retransmit pending_seg
				send_tcp_segment_via_ip(c->remote_ip, c->remote_port, c->local_port, c->pending_seg, c->pending_len);
				c->last_tx_time_ms = now;
				c->retransmit_count++;
			}
		}
	}
	// ARP tick etc. not here
}
